
`include "defines.v"

module axi_rw # (
    parameter RW_DATA_WIDTH     = 64,
    parameter RW_ADDR_WIDTH     = 64,
    parameter AXI_DATA_WIDTH    = 64,
    parameter AXI_ADDR_WIDTH    = 64,
    parameter AXI_ID_WIDTH      = 2,
    parameter AXI_USER_WIDTH    = 1,
    parameter AXI_ID            = 0
)(
    input  wire                              clock,
    input  wire                              reset_n,

	input                                    r_req,
    output wire                              r_handshaked,
	output reg                               r_okay,
    output reg  [RW_DATA_WIDTH-1:0]          data_read,
    input  wire [AXI_DATA_WIDTH-1:0]         r_addr,
    input  wire [1:0]                        r_size,
    output reg  [1:0]                        r_resp,

	input  wire                              w_req,
	output reg                               w_okay,
    input  wire [RW_DATA_WIDTH-1:0]          data_write,
    input  wire [AXI_DATA_WIDTH-1:0]         w_addr,
    input  wire [1:0]                        w_size,
    output reg  [1:0]                        w_resp,



    // Advanced eXtensible Interface
    input  wire                              axi_aw_ready_i,
    output wire                              axi_aw_valid_o,
    output wire [AXI_ADDR_WIDTH-1:0]         axi_aw_addr_o,
    output wire [2:0]                        axi_aw_prot_o,
    output wire [AXI_ID_WIDTH-1:0]           axi_aw_id_o,
    output wire [AXI_USER_WIDTH-1:0]         axi_aw_user_o,
    output wire [7:0]                        axi_aw_len_o,
    output wire [2:0]                        axi_aw_size_o,
    output wire [1:0]                        axi_aw_burst_o,
    output wire                              axi_aw_lock_o,
    output wire [3:0]                        axi_aw_cache_o,
    output wire [3:0]                        axi_aw_qos_o,
    output wire [3:0]                        axi_aw_region_o,

    input  wire                              axi_w_ready_i,
    output wire                              axi_w_valid_o,
    output reg  [AXI_DATA_WIDTH-1:0]         axi_w_data_o,
    output reg  [AXI_DATA_WIDTH/8-1:0]       axi_w_strb_o,
    output wire                              axi_w_last_o,
    output wire [AXI_USER_WIDTH-1:0]         axi_w_user_o,
    
    output wire                              axi_b_ready_o,
    input  wire                              axi_b_valid_i,
    input  wire [1:0]                        axi_b_resp_i,
    input  wire [AXI_ID_WIDTH-1:0]           axi_b_id_i,
    input  wire [AXI_USER_WIDTH-1:0]         axi_b_user_i,

    input  wire                              axi_ar_ready_i,
    output wire                              axi_ar_valid_o,
    output wire [AXI_ADDR_WIDTH-1:0]         axi_ar_addr_o,
    output wire [2:0]                        axi_ar_prot_o,
    output wire [AXI_ID_WIDTH-1:0]           axi_ar_id_o,
    output wire [AXI_USER_WIDTH-1:0]         axi_ar_user_o,
    output wire [7:0]                        axi_ar_len_o,
    output wire [2:0]                        axi_ar_size_o,
    output wire [1:0]                        axi_ar_burst_o,
    output wire                              axi_ar_lock_o,
    output wire [3:0]                        axi_ar_cache_o,
    output wire [3:0]                        axi_ar_qos_o,
    output wire [3:0]                        axi_ar_region_o,
    
    output wire                              axi_r_ready_o,
    input  wire                              axi_r_valid_i,
    input  wire [1:0]                        axi_r_resp_i,
    input  wire [AXI_DATA_WIDTH-1:0]         axi_r_data_i,
    input  wire                              axi_r_last_i,
    input  wire [AXI_ID_WIDTH-1:0]           axi_r_id_i,
    input  wire [AXI_USER_WIDTH-1:0]         axi_r_user_i
);

    // handshake
    wire aw_hs      = axi_aw_ready_i & axi_aw_valid_o;
    wire w_hs       = axi_w_ready_i  & axi_w_valid_o;
    wire b_hs       = axi_b_ready_o  & axi_b_valid_i;
    wire ar_hs      = axi_ar_ready_i & axi_ar_valid_o;
    wire r_hs       = axi_r_ready_o  & axi_r_valid_i;

    wire w_done     = w_hs & axi_w_last_o;
    wire r_done     = r_hs & axi_r_last_i;
    wire b_done     = b_hs;

    
    // ------------------State Machine------------------
    parameter [1:0] W_STATE_IDLE = 2'b00, W_STATE_ADDR = 2'b01, W_STATE_WRITE = 2'b10, W_STATE_RESP = 2'b11;
    parameter [1:0] R_STATE_IDLE = 2'b00, R_STATE_ADDR = 2'b01, R_STATE_READ  = 2'b10;

    reg [1:0] w_state, r_state;
    wire w_state_idle = w_state == W_STATE_IDLE, w_state_addr = w_state == W_STATE_ADDR, w_state_write = w_state == W_STATE_WRITE, w_state_resp = w_state == W_STATE_RESP;
    wire r_state_idle = r_state == R_STATE_IDLE, r_state_addr = r_state == R_STATE_ADDR, r_state_read  = r_state == R_STATE_READ;

    // Wirte State Machine
    always @(posedge clock) begin
        if (~reset_n) begin
            w_state <= W_STATE_IDLE;
        end
        else begin
            if (w_req) begin
                case (w_state)
                    W_STATE_IDLE:               w_state <= W_STATE_ADDR;
                    W_STATE_ADDR:  if (aw_hs)   w_state <= W_STATE_WRITE;
                    W_STATE_WRITE: if (w_done)  w_state <= W_STATE_RESP;
                    W_STATE_RESP:  if (b_hs)    w_state <= W_STATE_IDLE;
                    default : ;
                endcase
            end
        end
    end

    // Read State Machine
    always @(posedge clock) begin
        if (~reset_n) begin
            r_state <= R_STATE_IDLE;
        end
        else begin
            if (r_req) begin
                case (r_state)
                    R_STATE_IDLE:               r_state <= R_STATE_ADDR;
                    R_STATE_ADDR: if (ar_hs)    r_state <= R_STATE_READ;
                    R_STATE_READ: if (r_done)   r_state <= R_STATE_IDLE;
                    default:;
                endcase
            end
        end
    end


    // ------------------Number of transmission------------------
    //-------------------read----------------------
    reg [7:0] r_len;
    wire r_len_reset_n      = reset_n & ~(r_req & r_state_idle);
    wire r_len_incr_en    = (r_len != axi_r_len) & r_hs;
    always @(posedge clock) begin
        if (~r_len_reset_n) begin
            r_len <= 0;
        end
        else if (r_len_incr_en) begin
            r_len <= r_len + 1;
        end
    end

    //-------------------write----------------------
    reg [7:0] w_len;
    wire w_len_reset_n      = reset_n & ~(w_req & b_done);
    wire w_len_incr_en    = (w_len != axi_w_len) & w_hs;
    always @(posedge clock) begin
        if (~w_len_reset_n) begin
            w_len <= 0;
        end
        else if (w_len_incr_en) begin
            w_len <= w_len + 1;
        end
    end

    // ------------------Process Data------------------
    parameter ALIGNED_WIDTH = $clog2(AXI_DATA_WIDTH / 8);
    parameter OFFSET_WIDTH  = $clog2(AXI_DATA_WIDTH);
    parameter AXI_SIZE      = $clog2(AXI_DATA_WIDTH / 8);
    parameter MASK_WIDTH    = AXI_DATA_WIDTH * 2;
    parameter TRANS_LEN     = RW_DATA_WIDTH / AXI_DATA_WIDTH;
    parameter BLOCK_TRANS   = TRANS_LEN > 1 ? 1'b1 : 1'b0;

    //-------------------read----------------------
    wire r_aligned            = BLOCK_TRANS | r_addr[ALIGNED_WIDTH-1:0] == 0;
    wire r_size_b             = r_size == `SIZE_B;
    wire r_size_h             = r_size == `SIZE_H;
    wire r_size_w             = r_size == `SIZE_W;
    wire r_size_d             = r_size == `SIZE_D;
    wire [3:0] r_addr_op1     = {{4-ALIGNED_WIDTH{1'b0}}, r_addr[ALIGNED_WIDTH-1:0]};
    wire [3:0] r_addr_op2     = ({4{r_size_b}} & {4'b0})
                                | ({4{r_size_h}} & {4'b1})
                                | ({4{r_size_w}} & {4'b11})
                                | ({4{r_size_d}} & {4'b111})
                                ;
    wire [3:0] r_addr_end     = r_addr_op1 + r_addr_op2;
    wire r_overstep           = r_addr_end[3:ALIGNED_WIDTH] != 0;

    wire [7:0] axi_r_len      = r_aligned ? TRANS_LEN - 1 : {{7{1'b0}}, r_overstep};
    wire [2:0] axi_r_size     = AXI_SIZE[2:0];
    
    wire [AXI_ADDR_WIDTH-1:0] axi_r_addr          = {r_addr[AXI_ADDR_WIDTH-1:ALIGNED_WIDTH], {ALIGNED_WIDTH{1'b0}}};
    wire [OFFSET_WIDTH-1:0] r_aligned_offset_l    = {{OFFSET_WIDTH-ALIGNED_WIDTH{1'b0}}, {r_addr[ALIGNED_WIDTH-1:0]}} << 3;
    wire [OFFSET_WIDTH-1:0] r_aligned_offset_h    = ~r_aligned_offset_l + 1'b1;
    wire [MASK_WIDTH-1:0] r_mask                  = (({MASK_WIDTH{r_size_b}} & {{MASK_WIDTH-8{1'b0}}, 8'hff})
                                                    | ({MASK_WIDTH{r_size_h}} & {{MASK_WIDTH-16{1'b0}}, 16'hffff})
                                                    | ({MASK_WIDTH{r_size_w}} & {{MASK_WIDTH-32{1'b0}}, 32'hffffffff})
                                                    | ({MASK_WIDTH{r_size_d}} & {{MASK_WIDTH-64{1'b0}}, 64'hffffffff_ffffffff})
                                                    ) << r_aligned_offset_l;
    wire [AXI_DATA_WIDTH-1:0] r_mask_l            = r_mask[AXI_DATA_WIDTH-1:0];
    wire [AXI_DATA_WIDTH-1:0] r_mask_h            = r_mask[MASK_WIDTH-1:AXI_DATA_WIDTH];

    wire [AXI_ID_WIDTH-1:0] axi_r_id              = AXI_ID;
    wire [AXI_USER_WIDTH-1:0] axi_r_user          =  {AXI_USER_WIDTH{1'b0}};

    wire r_okay_nxt = r_done;
    wire r_okay_en      = r_done | r_okay;
    always @(posedge clock) begin
        if (~reset_n) begin
            r_okay <= 0;
        end
        else if (r_okay_en) begin
            r_okay <= r_okay_nxt;
        end
    end

    wire [1:0] r_resp_nxt = axi_r_resp_i;
    wire r_resp_en = r_done;
    always @(posedge clock) begin
        if (~reset_n) begin
            r_resp <= 0;
        end
        else if (r_resp_en) begin
            r_resp <= r_resp_nxt;
        end
    end
    assign r_handshaked = r_state_read | (ar_hs & r_req);

    //-------------------write----------------------
    wire w_aligned            = BLOCK_TRANS | r_addr[ALIGNED_WIDTH-1:0] == 0;
    wire w_size_b             = w_size == `SIZE_B;
    wire w_size_h             = w_size == `SIZE_H;
    wire w_size_w             = w_size == `SIZE_W;
    wire w_size_d             = w_size == `SIZE_D;
    wire [3:0] w_addr_op1     = {{4-ALIGNED_WIDTH{1'b0}}, w_addr[ALIGNED_WIDTH-1:0]};
    wire [3:0] w_addr_op2     = ({4{w_size_b}} & {4'b0})
                                | ({4{w_size_h}} & {4'b1})
                                | ({4{w_size_w}} & {4'b11})
                                | ({4{w_size_d}} & {4'b111})
                                ;
    wire [3:0] w_addr_end     = w_addr_op1 + w_addr_op2;
    wire w_overstep           = w_addr_end[3:ALIGNED_WIDTH] != 0;

    wire [7:0] axi_w_len      = w_aligned ? TRANS_LEN - 1 : {{7{1'b0}}, w_overstep};
    wire [2:0] axi_w_size     = AXI_SIZE[2:0];
    
    wire [AXI_ADDR_WIDTH-1:0] axi_w_addr          = {w_addr[AXI_ADDR_WIDTH-1:ALIGNED_WIDTH], {ALIGNED_WIDTH{1'b0}}};
    wire [OFFSET_WIDTH-1:0] w_aligned_offset_l    = {{OFFSET_WIDTH-ALIGNED_WIDTH{1'b0}}, {w_addr[ALIGNED_WIDTH-1:0]}} << 3;
    wire [OFFSET_WIDTH-1:0] w_aligned_offset_h    = ~w_aligned_offset_l + 1'b1;
    wire [AXI_DATA_WIDTH/4-1:0] w_mask            = (({AXI_DATA_WIDTH/4{w_size_b}} & {{AXI_DATA_WIDTH/4-1{1'b0}}, 1'b1})
                                                    | ({AXI_DATA_WIDTH/4{w_size_h}} & {{AXI_DATA_WIDTH/4-2{1'b0}}, 2'b11})
                                                    | ({AXI_DATA_WIDTH/4{w_size_w}} & {{AXI_DATA_WIDTH/4-4{1'b0}}, 4'b1111})
                                                    | ({AXI_DATA_WIDTH/4{w_size_d}} & {{AXI_DATA_WIDTH/4-8{1'b0}}, 8'b1111_1111})
                                                    ) << {{OFFSET_WIDTH-ALIGNED_WIDTH{1'b0}}, {w_addr[ALIGNED_WIDTH-1:0]}};
    wire [AXI_DATA_WIDTH/8-1:0] w_strb_l = w_mask[AXI_DATA_WIDTH/8-1:0];
    wire [AXI_DATA_WIDTH/8-1:0] w_strb_h = w_mask[AXI_DATA_WIDTH/4-1:AXI_DATA_WIDTH/8];

    wire [AXI_ID_WIDTH-1:0] axi_w_id              = AXI_ID;
    wire [AXI_USER_WIDTH-1:0] axi_w_user          =  {AXI_USER_WIDTH{1'b0}};

    wire w_okay_nxt = b_done;
    wire w_okay_en      = b_done | w_okay;
    always @(posedge clock) begin
        if (~reset_n) begin
            w_okay <= 0;
        end
        else if (w_okay_en) begin
            w_okay <= w_okay_nxt;
        end
    end

    wire [1:0] w_resp_nxt = axi_b_resp_i;
    wire w_resp_en = b_done;
    always @(posedge clock) begin
        if (~reset_n) begin
            w_resp <= 0;
        end
        else if (w_resp_en) begin
            w_resp <= w_resp_nxt;
        end
    end


    // ------------------Write Transaction------------------
    // Write address channel signals
    assign axi_aw_valid_o = w_state_addr;
    assign axi_aw_addr_o  = axi_w_addr;
    assign axi_aw_prot_o  = `AXI_PROT_UNPRIVILEGED_ACCESS | `AXI_PROT_SECURE_ACCESS | `AXI_PROT_DATA_ACCESS;
    assign axi_aw_id_o    = axi_w_id;
    assign axi_aw_user_o  = axi_w_user;
    assign axi_aw_len_o   = axi_w_len;
    assign axi_aw_size_o  = axi_w_size;
    assign axi_aw_burst_o = `AXI_BURST_TYPE_INCR;
    assign axi_aw_lock_o  = 1'b0;
    assign axi_aw_cache_o = `AXI_ARCACHE_NORMAL_NON_CACHEABLE_NON_BUFFERABLE;
    assign axi_aw_qos_o   = 4'h0;

    // Write data channel signals
    assign axi_w_valid_o    = w_state_write;
    assign axi_w_last_o     = (axi_w_valid_o & axi_w_len == w_len);

    //Write response channel signals
    assign axi_b_ready_o    = w_state_resp;

    wire [AXI_DATA_WIDTH-1:0] axi_w_data_l  = data_write << w_aligned_offset_l;
    wire [AXI_DATA_WIDTH-1:0] axi_w_data_h  = data_write >> w_aligned_offset_h;

     always @(posedge clock) begin
         if (~reset_n) begin
             axi_w_data_o[AXI_DATA_WIDTH-1:0] <= 0;
             axi_w_strb_o <= 0;
         end
         else if (aw_hs | axi_w_ready_i & axi_w_valid_o) begin
             if (~w_aligned | w_overstep) begin
                 if (w_len[0]) begin
                     axi_w_data_o <= axi_w_data_h;
                     axi_w_strb_o <= w_strb_h;
                 end
                 else begin
                     axi_w_data_o <= axi_w_data_l;
                     axi_w_strb_o <= w_strb_l;
                 end
             end
             else if (w_len == 0) begin
                 axi_w_data_o[AXI_DATA_WIDTH-1:0] <= axi_w_data_l;
                 axi_w_strb_o <= w_strb_l;
             end
         end
     end


    // ------------------Read Transaction------------------

    // Read address channel signals
    assign axi_ar_valid_o   = r_state_addr;
    assign axi_ar_addr_o    = axi_r_addr;
    assign axi_ar_prot_o    = `AXI_PROT_UNPRIVILEGED_ACCESS | `AXI_PROT_SECURE_ACCESS | `AXI_PROT_DATA_ACCESS;
    assign axi_ar_id_o      = axi_r_id;
    assign axi_ar_user_o    = axi_r_user;
    assign axi_ar_len_o     = axi_r_len;
    assign axi_ar_size_o    = axi_r_size;
    assign axi_ar_burst_o   = `AXI_BURST_TYPE_INCR;
    assign axi_ar_lock_o    = 1'b0;
    assign axi_ar_cache_o   = `AXI_ARCACHE_NORMAL_NON_CACHEABLE_NON_BUFFERABLE;
    assign axi_ar_qos_o     = 4'h0;

    // Read data channel signals
    assign axi_r_ready_o    = r_state_read;

    wire [AXI_DATA_WIDTH-1:0] axi_r_data_l  = (axi_r_data_i & r_mask_l) >> r_aligned_offset_l;
    wire [AXI_DATA_WIDTH-1:0] axi_r_data_h  = (axi_r_data_i & r_mask_h) << r_aligned_offset_h;


    always @(posedge clock) begin
        if (~reset_n) begin
            data_read[AXI_DATA_WIDTH-1:0] <= 0;
        end
        else if (axi_r_ready_o & axi_r_valid_i) begin
            if (~r_aligned | r_overstep) begin
                if (r_len[0]) begin
                    data_read[AXI_DATA_WIDTH-1:0] <= data_read[AXI_DATA_WIDTH-1:0] | axi_r_data_h;
                end
                else begin
                    data_read[AXI_DATA_WIDTH-1:0] <= axi_r_data_l;
                end
            end
            else if (r_len == 0) begin
                data_read[AXI_DATA_WIDTH-1:0] <= axi_r_data_l;
            end
        end
    end

endmodule
